import chromadb
import os

chroma_client = chromadb.PersistentClient(path='./mychroma_db')
chroma_client.heartbeat()


# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = "sk-w0P1gNv42gN6suR8819858Ee297c4b7cBc97Ae3e7b8dA701"
os.environ["OPENAI_API_BASE"] = "https://ai-yyds.com/v1"

from chromadb.utils import embedding_functions
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                model_name="text-embedding-ada-002",
                api_key = "sk-w0P1gNv42gN6suR8819858Ee297c4b7cBc97Ae3e7b8dA701")

collection = chroma_client.get_or_create_collection(name='deptName', embedding_function=openai_ef)